"""

"""
import os

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import CharacterTextSplitter
from my_huggingface.ModelScopeEmbeddings import ModelScopeEmbeddings

# 解决Huggingface 警告
os.environ['TOKENIZERS_PARALLELISM'] = "False"

loader = TextLoader("knowledge1.txt", encoding="UTF-8")
documents = loader.load()

text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# embedding model
local_model_path = "/Users/brightzhou/.cache/modelscope/hub/models/sentence-transformers/all-MiniLM-L6-v2"
embedding = ModelScopeEmbeddings(local_model_path, device='cpu')

## db
db = FAISS.from_documents(docs, embedding=embedding)
query = "Pixar公司是做什么的"

retriever = db.as_retriever()
result = retriever.invoke(query)
print(result[0].page_content)

print("--" * 50)
result_core_list = db.similarity_search_with_relevance_scores(query, k=1)
print(result_core_list)
for doc, score in result_core_list:
    print(score)
